import os
import numpy as np

import random
from tqdm import tqdm
import json
import matplotlib.pyplot as plt


def read_split_data(data_root, save_dir, val_rate=0.2):
    assert os.path.exists(data_root), "data path:{} does not exists".format(data_root)

    # 遍历文件夹，一个文件夹对应一个类别
    classes = [cla for cla in os.listdir(data_root) if os.path.isdir(os.path.join(data_root, cla))]
    # 排序，保证顺序一致
    classes.sort()

    def get_label(tmp_subfolder):
        label_dic = {'HCC': 3, 'CYST': 0, 'FNH': 1, 'HA': 2, 'ICC': 4, 'META': 5, 'Hemangioma': 2, 'nodule': 0}
        if 'HCC' in tmp_subfolder:
            return label_dic['HCC']
        elif 'CYST' in tmp_subfolder:
            return label_dic['CYST']
        elif 'FNH' in tmp_subfolder:
            return label_dic['FNH']
        elif 'HA' in tmp_subfolder:
            return label_dic['HA']
        elif 'ICC' in tmp_subfolder:
            return label_dic['ICC']
        elif 'META' in tmp_subfolder:
            return label_dic['META']
        elif 'Hemangioma' in tmp_subfolder:
            return label_dic['Hemangioma']
        elif 'nodule' in tmp_subfolder:
            return label_dic['nodule']

    images_path = []
    train_images_path = []  # 存储训练集的所有图片路径
    train_images_label = []  # 存储训练集图片对应索引信息
    val_images_path = []  # 存储验证集的所有图片路径
    val_images_label = []  # 存储验证集图片对应索引信息
    every_class_num = []  # 存储每个类别的样本总数
    train_images_pth = '/home/hzt/whh/pycharm_project_179/runs/Apr03_17-31-30_ubuntu-SYS-7049GP-TRT/train.txt'
    val_images_pth = '/home/hzt/whh/pycharm_project_179/runs/Apr03_17-31-30_ubuntu-SYS-7049GP-TRT/val.txt'
    # train_txt = open(os.path.join(save_dir, 'train.txt'), 'w')
    # val_txt = open(os.path.join(save_dir, 'val.txt'), 'w')
    with open(train_images_pth, 'r') as file:
        # 逐行读取文件内容
        for line in file:
            train_images_path.append(line)
            train_images_path = [line.rstrip('\n') for line in train_images_path]

    num_train_images = len(train_images_path)
    for i in range(num_train_images):
        train_images_label.append(get_label(train_images_path[i]))

    with open(val_images_pth, 'r') as file:
        # 逐行读取文件内容
        for line in file:
            val_images_path.append(line)
            val_images_path = [line.rstrip('\n') for line in val_images_path]

    num_val_images = len(val_images_path)
    for i in range(num_val_images):
        val_images_label.append(get_label(val_images_path[i]))
    # 遍历每个文件夹下的文件
    # for cla in tqdm(classes):
    #     cla_path = os.path.join(data_root, cla)
    #     images = os.listdir(cla_path)  # 获取该类别下的所有文件名
    #     if '.DS_Store' in images:
    #         images.remove('.DS_Store')
    #     for patient in images:
    #         patient_path = os.path.join(cla_path,patient)
    #         images_path.append(patient_path)
    # random.shuffle(images_path)
    # num_images = len(images_path)  # 获取该类别下的总样本数
    # num_val = int(num_images * val_rate)  # 计算验证集的样本数
    # for i in range(num_images):
    #     if i < num_val:
    #         val_txt.write(images_path[i] + "\n")
    #         val_images_path.append(images_path[i])
    #         val_images_label.append(get_label(images_path[i]))
    #     else:
    #         train_txt.write(images_path[i] + "\n")  # 写入训练集文件路径
    #         train_images_path.append(images_path[i])  # 添加到训练集路径列表
    #         train_images_label.append(get_label(images_path[i]))  # 添加对应标签到训练集标签列表
    every_class_num.append(len(set(classes)))  # 记录该类别的总样本数
    # train_txt.close()
    # val_txt.close()

    print("{} images were found in the dataset.".format(len(train_images_path)+len(val_images_path)))
    print("{} images for training.".format(len(train_images_path)))
    print("{} images for validation.".format(len(val_images_path)))

    return train_images_path, val_images_path, train_images_label, val_images_label, every_class_num


if __name__ == '__main__':
    # read_split_data(data_root="/Users/hern/Desktop/毕设/Data", save_dir="/Users/hern/Desktop/毕设/Data")
    #计算随机分划的肿瘤类别分别有多少个数据
    train_images_pth = '/home/hzt/whh/pycharm_project_179/runs/Apr03_17-31-30_ubuntu-SYS-7049GP-TRT/train.txt'
    val_images_pth = '/home/hzt/whh/pycharm_project_179/runs/Apr03_17-31-30_ubuntu-SYS-7049GP-TRT/val.txt'
    def get_label(tmp_subfolder):
        label_dic = {'HCC': 3, 'CYST': 0, 'FNH': 1, 'HA': 2, 'ICC': 4, 'META': 5, 'Hemangioma': 2, 'nodule': 0}
        if 'HCC' in tmp_subfolder:
            return label_dic['HCC']
        elif 'CYST' in tmp_subfolder:
            return label_dic['CYST']
        elif 'FNH' in tmp_subfolder:
            return label_dic['FNH']
        elif 'HA' in tmp_subfolder:
            return label_dic['HA']
        elif 'ICC' in tmp_subfolder:
            return label_dic['ICC']
        elif 'META' in tmp_subfolder:
            return label_dic['META']
        elif 'Hemangioma' in tmp_subfolder:
            return label_dic['Hemangioma']
        elif 'nodule' in tmp_subfolder:
            return label_dic['nodule']


    train_images_path = []  # 存储训练集的所有图片路径
    train_images_label = []  # 存储训练集图片对应索引信息
    val_images_path = []  # 存储验证集的所有图片路径
    val_images_label = []
    with open(train_images_pth, 'r') as file:
        # 逐行读取文件内容
        for line in file:
            train_images_path.append(line)
            train_images_path = [line.rstrip('\n') for line in train_images_path]

    num_train_images = len(train_images_path)
    for i in range(num_train_images):
        train_images_label.append(get_label(train_images_path[i]))

    with open(val_images_pth, 'r') as file:
        # 逐行读取文件内容
        for line in file:
            val_images_path.append(line)
            val_images_path = [line.rstrip('\n') for line in val_images_path]

    num_val_images = len(val_images_path)
    for i in range(num_val_images):
        val_images_label.append(get_label(val_images_path[i]))

    for i in range(6):
        print("train_",i,"=",train_images_label.count(i))
    for i in range(6):
        print("val_",i,"=",val_images_label.count(i))
